package com.gitee.helioz.openai.api;

import com.gitee.helioz.openai.common.error.OpenAiRuntimeException;
import com.gitee.helioz.openai.common.util.http.ContentTypeEnum;
import com.gitee.helioz.openai.common.util.http.HttpFilePart;
import com.gitee.helioz.openai.common.util.http.HttpProxyConfig;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.asynchttpclient.AsyncHttpClient;
import org.asynchttpclient.DefaultAsyncHttpClient;
import org.asynchttpclient.DefaultAsyncHttpClientConfig;
import org.asynchttpclient.Realm;
import org.asynchttpclient.Request;
import org.asynchttpclient.RequestBuilder;
import org.asynchttpclient.proxy.ProxyServer;
import org.asynchttpclient.proxy.ProxyType;
import org.asynchttpclient.request.body.multipart.ByteArrayPart;
import org.asynchttpclient.request.body.multipart.StringPart;

import java.io.IOException;
import java.net.Proxy;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutionException;

/**
 * 〈AHC实现〉
 *
 * @author <a href="https://gitee.com/helioz">Morning Starr</a>
 */
@Slf4j
public class OpenAiAHCServiceImpl extends BaseOpenAiServiceImpl {

    private AsyncHttpClient asyncHttpClient;

    @Override
    protected void initHttp() {
        String threadPoolName = "AHC-THREAD-POOL";
        int ioThreadsCount = Runtime.getRuntime().availableProcessors() * 2;
        log.debug("ioThreadsCount={}", ioThreadsCount);
        // 设置代理
        DefaultAsyncHttpClientConfig.Builder ahcConfigBuilder = new DefaultAsyncHttpClientConfig.Builder();
        ahcConfigBuilder.setThreadPoolName(threadPoolName).setIoThreadsCount(ioThreadsCount);
        if (this.getHttpProxyConfig() != null) {
            HttpProxyConfig httpProxyConfig = this.getHttpProxyConfig();
            ProxyServer.Builder proxyBuilder = new ProxyServer.Builder(httpProxyConfig.getHost(), httpProxyConfig.getPort());
            if (Proxy.Type.SOCKS.equals(httpProxyConfig.getType())) {
                proxyBuilder.setProxyType(ProxyType.SOCKS_V5);
            }
            ahcConfigBuilder.setProxyServer(proxyBuilder);

            // 设置代理认证账号密码
            if (!StringUtils.isAnyBlank(httpProxyConfig.getName(), httpProxyConfig.getPassword())) {
                ahcConfigBuilder.setRealm(new Realm.Builder(httpProxyConfig.getName(), httpProxyConfig.getPassword()).build());
            }
        }
        this.asyncHttpClient = new DefaultAsyncHttpClient(ahcConfigBuilder.build());
        log.info("初始化AHC完成!");
    }

    @Override
    protected String sendGetRequest(String url) throws IOException {
        Request request = new RequestBuilder("GET")
                .setUrl(url)
                .setSingleHeaders(buildHeaderMap())
                .build();
        String responseBody;
        try {
            responseBody = asyncHttpClient.executeRequest(request).get().getResponseBody();
        } catch (InterruptedException | ExecutionException ex) {
            throw new OpenAiRuntimeException(ex);
        }
        return responseBody;
    }

    @Override
    protected String sendPostJsonRequest(String url, String jsonBody) throws IOException {
        Request request = new RequestBuilder("POST")
                .setUrl(url)
                .setSingleHeaders(buildHeaderMap())
                .addHeader("Content-Type", ContentTypeEnum.JSON)
                .setBody(jsonBody)
                .build();
        String responseBody;
        try {
            responseBody = asyncHttpClient.executeRequest(request).get().getResponseBody();
        } catch (InterruptedException | ExecutionException ex) {
            throw new OpenAiRuntimeException(ex);
        }
        return responseBody;
    }

    @Override
    protected String sendPostFormData(String url, Map<String, Object> formMap, List<HttpFilePart> fileParts) throws IOException {
        RequestBuilder requestBuilder = new RequestBuilder("POST");
        requestBuilder.setUrl(url);
        requestBuilder.setSingleHeaders(buildHeaderMap());
        if (fileParts != null && !fileParts.isEmpty()) {
            requestBuilder.addHeader("Content-Type", ContentTypeEnum.MULTIPART);
            fileParts.forEach(filePart -> requestBuilder.addBodyPart(new ByteArrayPart(filePart.getPartName(), filePart.getFileBytes())));
            for (Map.Entry<String, Object> formItem : formMap.entrySet()) {
                if (formItem.getValue() != null) {
                    requestBuilder.addBodyPart(new StringPart(formItem.getKey(), String.valueOf(formItem.getValue())));
                }
            }
        } else {
            requestBuilder.addHeader("Content-Type", ContentTypeEnum.FORM_URLENCODED);
            for (Map.Entry<String, Object> formItem : formMap.entrySet()) {
                if (formItem.getValue() != null) {
                    requestBuilder.addFormParam(formItem.getKey(), String.valueOf(formItem.getValue()));
                }
            }
        }
        String responseBody;
        try {
            responseBody = asyncHttpClient.executeRequest(requestBuilder.build()).get().getResponseBody();
        } catch (InterruptedException | ExecutionException ex) {
            throw new OpenAiRuntimeException(ex);
        }
        return responseBody;
    }

    @Override
    protected String sendDeleteRequest(String url) throws IOException {
        Request request = new RequestBuilder("DELETE")
                .setUrl(url)
                .build();
        String responseBody;
        try {
            responseBody = asyncHttpClient.executeRequest(request).get().getResponseBody();
        } catch (InterruptedException | ExecutionException ex) {
            throw new OpenAiRuntimeException(ex);
        }
        return responseBody;
    }

    public void close() {
        try {
            asyncHttpClient.close();
        } catch (IOException ioe) {
            log.warn("关闭AHC客户端发生异常, error={}", ioe.getMessage(), ioe);
        }
    }

    @Override
    public void setHttpProxyConfig(HttpProxyConfig httpProxyConfig) {
        // 变更HTTP代理配置
        if (this.getHttpProxyConfig() == null) {
            this.asyncHttpClient = new DefaultAsyncHttpClient();
        }
    }

}
